#!/usr/bin/env python
# -*- coding: utf-8 -*-
# datetime: 31/3/2024 上午11:56
# Author: name
# File: custom_data_extract.py
# ide: PyCharm
# description: 数据提取工具类
import re
import time
import traceback
from typing import Union

import jmespath
import requests
from requests import Response


class Extract:
    json: Union[list] = None
    re: Union[list] = None

    def __init__(self, **entries):
        self.__dict__.update(entries)


class DataExtract:
    def __init__(self, actual_response: Response, extract_data: Extract):
        self.response = actual_response
        self.extract = extract_data
        self.variables = {}
        self.log_list = []

    @staticmethod
    def is_valid_jmespath(expression, log_list=None):
        """
            校验jmespath表达式是否有效
        :param log_list:
        :param expression:
        :return:
        """
        if log_list is None:
            log_list = []
        try:
            jmespath.compile(expression)
            return True
        except Exception:
            log_list.append(traceback.format_exc())
            return False

    @staticmethod
    def is_valid_regex(regex_str):
        """
            校验正则表达式是否有效
        :param regex_str:
        :return:
        """
        try:
            re.compile(regex_str)
            return True
        except re.error:
            return False

    def extract_json(self):
        """
            json格式数据提取,将提取的数据保存到self.variables对象中
        :return:
        """
        if "application/json" in response.headers.get("content-type"):
            if isinstance(self.extract.json, list):
                for item in self.extract.json:
                    print(item)
                    if isinstance(item, dict):
                        if self.is_valid_jmespath(item.get("path"), self.log_list):
                            self.variables[item.get("name")] = jmespath.search(item.get("path"), self.response.json())
                            self.log_list.append(
                                f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取成功,变量: {item.get("name")}={self.variables[item.get("name")]}')
                        else:
                            self.log_list.append(
                                f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取失败,变量: {item.get("name")}对应的jmespath: {item.get("path")}存在语法错误')
                    elif isinstance(item, str):
                        if "=" in item:
                            key = item.split("=")[0]
                            expression = item[len(key) + 1:]
                            if self.is_valid_jmespath(expression, self.log_list):
                                self.variables[key] = jmespath.search(expression, self.response.json())
                                self.log_list.append(
                                    f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取成功,变量{key}={self.variables[key]}')
                            else:
                                self.log_list.append(
                                    f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取失败,变量: {key}对应的jmespath: {expression}存在语法错误')
                        else:
                            self.log_list.append(
                                f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取失败,jmespath: {item}存在语法错误')
                    else:
                        pass
            elif isinstance(self.extract.json, str):
                item = self.extract.json
                if "=" in item:
                    key = item.split("=")[0]
                    expression = item[len(key) + 1:]
                    if self.is_valid_jmespath(expression, self.log_list):
                        self.variables[key] = jmespath.search(expression, self.response.json())
                        self.log_list.append(
                            f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取成功,变量{key}={self.variables[key]}')
                    else:
                        self.log_list.append(
                            f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取失败,变量: {key}对应的jmespath: {expression}存在语法错误')
                else:
                    self.log_list.append(
                        f'{time.strftime("%Y-%m-%d %H:%M:%S")}_json数据提取失败,jmespath: {item}存在语法错误')

    def extract_re(self):
        if isinstance(self.extract.re, list):
            for item in self.extract.re:
                if isinstance(item, dict):
                    if self.is_valid_regex(item.get("expression")):
                        match = re.findall(item.get("expression"), self.response.text)
                        if match:
                            self.variables[item.get('name')] = match[int(item.get("index"))]
                            self.log_list.append(
                                f"{time.strftime('%Y-%m-%d %H:%M:%S')}_正则表达式变量提取成功: 变量:{item.get('name')}={self.variables[item.get('name')]}")
                        else:
                            self.log_list.append(
                                f"{time.strftime('%Y-%m-%d %H:%M:%S')}_正则表达式变量提取失败: 变量{item.get('name')}匹配结果为:None")
                elif isinstance(item, str):
                    if "=" in item:
                        key = item.split("=")[0]
                        expression = item[len(key) + 1:]
                        if self.is_valid_regex(expression):
                            match = re.findall(expression, self.response.text)
                            if match:
                                self.variables[key] = match[0]
                                self.log_list.append(
                                    f"{time.strftime('%Y-%m-%d %H:%M:%S')}_正则表达式变量提取成功: 变量:{key}={self.variables[key]}")
                            else:
                                self.log_list.append(
                                    f"{time.strftime('%Y-%m-%d %H:%M:%S')}_正则表达式变量提取失败: 变量{key}匹配结果为:None")

        elif isinstance(self.extract.re, str):
            if "=" in self.extract.re:
                key = self.extract.re.split("=")[0]
                expression = self.extract.re[len(key) + 1:]
                if self.is_valid_regex(expression):
                    match = re.findall(expression, self.response.text)
                    if match:
                        self.variables[key] = match[0]
                        self.log_list.append(
                            f"{time.strftime('%Y-%m-%d %H:%M:%S')}_正则表达式变量提取成功: 变量:{key}={self.variables[key]}")
                    else:
                        self.log_list.append(
                            f"{time.strftime('%Y-%m-%d %H:%M:%S')}_正则表达式变量提取失败: 变量{key}匹配结果为:None")

        else:
            pass


if __name__ == '__main__':
    extract = {
        "json": [
            {
                "name": "page_size",
                "path": "current_page"
            },
            "total=total",
            "page_size=page_size",
            "address=results[0].address"

        ],
        "re": [
            {
                "name": "row_id",
                "expression": "\"id\":(.*?),",
                "index": "0"
            }
            ,
            'username="username":"(.*?)",',
            'username2="username2":"(.*?)",'
        ]
    }
    extract = {
        "json": "total=total",
        "re": 'username="username":"(.*?)",'
    }
    response = requests.get("http://127.0.0.1:8081/server/page?page=1&size=10&name=")
    print(response.text)
    de = DataExtract(response, Extract(**extract))
    de.extract_json()
    de.extract_re()
    print(de.variables)
    print(de.log_list)
